import pygame
import sys
import random
import math
from pygame.locals import *


# define game parameters
FPS = 30
WINDOW_WIDTH = 600
WINDOW_HEIGHT = 600
BOARD_RADIUS = 150  # board radius is 150 px

# define color constants
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)


# AI
def ai_move(pieces):
    # get all movable pieces
    movable_pieces = []
    for i in range(len(pieces)):
        if pieces[i] is not None:
            moves = get_available_moves(i, pieces)
            if len(moves) > 0:
                movable_pieces.append((i, moves))
    if len(movable_pieces) == 0:
        return None, None
    # randomly choose a movable piece
    piece, moves = random.choice(movable_pieces)
    # randomly choose a move
    move = random.choice(moves)
    return piece, move


def get_available_moves(piece, pieces):
    row, col = get_piece_coord(piece)
    moves = []
    for d_row, d_col in [(0, 2), (0, -2), (2, 0), (-2, 0), (2, 2), (-2, -2)]:
        jumped_piece = get_piece_at_coord(row+d_row//2, col+d_col//2, pieces)
        landing_piece = get_piece_at_coord(row+d_row, col+d_col, pieces)
        if jumped_piece is not None and landing_piece is None:
            moves.append((piece, jumped_piece, landing_piece))
    return moves


def get_piece_coord(piece):
    row = piece // 5
    col = piece % 5
    return row, col


def get_piece_at_coord(row, col, pieces):
    if row < 0 or row >= 5 or col < 0 or col >= 5:
        return None
    return pieces[row*5 + col]

def get_piece_center(piece_index):
    """
    Get the center position of a piece, based on its index
    :param piece_index: integer, index of the current piece
    :return: tuple, center position of the piece
    """
    row, col = get_piece_coord(piece_index)
    x = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.cos(angle_step*(5*row+col)/5))
    y = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.sin(angle_step*(5*row+col)/5))
    return (x, y)

def get_clicked_piece(x, y):
    """
    Get the index of the piece that the player clicked on
    :param x: integer, x coordinate of the click
    :param y: integer, y coordinate of the click
    :return: integer, index of the clicked piece, or None if no piece is clicked
    """
    for i in range(len(pieces)):
        if pieces[i] is not None:
            row, col = get_piece_coord(i)
            x0 = int(BOARD_RADIUS + (BOARD_RADIUS - 3 * space) * math.cos(angle_step * (5 * row + col) / 5))
            y0 = int(BOARD_RADIUS + (BOARD_RADIUS - 3 * space) * math.sin(angle_step * (5 * row + col) / 5))
            if (x - x0) ** 2 + (y - y0) ** 2 <= 10 ** 2:
                return i
    return None

# initialize pygame
pygame.init()
fps_clock = pygame.time.Clock()


# set window and board
window = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
pygame.display.set_caption('Jumping Game')

board = pygame.Surface((BOARD_RADIUS*2, BOARD_RADIUS*2))
board.fill(WHITE)
board.set_colorkey(WHITE)

points = [(BOARD_RADIUS, 0), (0, BOARD_RADIUS), (BOARD_RADIUS, BOARD_RADIUS*2),
          (BOARD_RADIUS*2, BOARD_RADIUS), (BOARD_RADIUS, 0)]
pygame.draw.polygon(board, BLACK, points, 3)

positions = [(BOARD_RADIUS, BOARD_RADIUS), (BOARD_RADIUS, BOARD_RADIUS//3),
             (2*BOARD_RADIUS//3, BOARD_RADIUS//3), (BOARD_RADIUS//3, BOARD_RADIUS*2//3),
             (1*BOARD_RADIUS//3, BOARD_RADIUS)]
for pos in positions:
    pygame.draw.circle(board, BLACK, pos, 10)


window.blit(board, (WINDOW_WIDTH//2-BOARD_RADIUS, WINDOW_HEIGHT//2-BOARD_RADIUS))



# initialize pieces
pieces = []
space = 5  # distance between line and piece
angle_step = 2 * 3.14159265 / 10
for i in range(10):
    angle = angle_step * i
    x = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.cos(angle))
    y = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.sin(angle))
    if (x, y) in [(BOARD_RADIUS, BOARD_RADIUS), (BOARD_RADIUS, BOARD_RADIUS-space-2),
                  (BOARD_RADIUS-space-2, BOARD_RADIUS-space-2), (BOARD_RADIUS+space+2, BOARD_RADIUS-space-2),
                  (BOARD_RADIUS-space-2, BOARD_RADIUS+space+2)]:
        piece = pygame.Surface((20, 20))
        piece.fill(RED)
        piece.set_colorkey(WHITE)
        pygame.draw.circle(piece, BLACK, (10, 10), 10, 2)
    else:
        piece = None
    pieces.append(piece)



# let player select the first piece
pygame.display.set_caption('Jumping Game: Select a piece')
selected_piece = None
while selected_piece is None:
    for event in pygame.event.get():
        if event.type == QUIT:
            pygame.quit()
            sys.exit()
        elif event.type == MOUSEBUTTONDOWN:
            x, y = event.pos
            clicked_piece = get_clicked_piece(x, y)
            if clicked_piece is not None and clicked_piece in players_pieces:
                selected_piece = clicked_piece

    window.fill(WHITE)

    if selected_piece is not None:
        # draw pieces
        for i in range(len(pieces)):
            if pieces[i] is not None:
                row, col = get_piece_coord(i)
                x = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.cos(angle_step*(5*row+col)/5))
                y = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.sin(angle_step*(5*row+col)/5))
                rect = pygame.Rect(x-10, y-10, 20, 20)
                window.blit(pieces[i], rect)

        # mark selected piece with a yellow circle
        pygame.draw.circle(window, (255, 255, 0), get_piece_center(selected_piece), 15, 3)

    # update display
    pygame.display.update()
    fps_clock.tick(FPS)


# game loop
player_turn = True
while True:
    for event in pygame.event.get():
        if event.type == QUIT:
            pygame.quit()
            sys.exit()

    if player_turn:
        for i in range(len(pieces)):
            if pieces[i] is not None and i != selected_piece:
                piece_x, piece_y = get_piece_coord(i)
                piece_x = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.cos(angle_step*(5*piece_x+piece_y)/5))
                piece_y = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.sin(angle_step*(5*piece_x+piece_y)/5))
                piece_rect = pygame.Surface((20, 20)).get_rect(center=(piece_x, piece_y))
                pygame.draw.circle(window, RED, piece_rect.center, 10)
                if selected_piece is not None:
                    moves = get_available_moves(selected_piece, pieces)
                    for move in moves:
                        jumper, target, destination = move
                        if i == jumper:
                            dest_x, dest_y = get_piece_coord(destination)
                            dest_x = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.cos(angle_step*(5*dest_x+dest_y)/5))
                            dest_y = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.sin(angle_step*(5*dest_x+dest_y)/5))
                            pygame.draw.line(window, RED, piece_rect.center, (dest_x, dest_y), 2)

        if selected_piece is not None:
            pygame.draw.circle(window, RED, piece_rect.center, 10, 2)

        if event.type == MOUSEBUTTONUP:
            x, y = event.pos
            dest_row, dest_col = None, None
            for i in range(len(pieces)):
                if pieces[i] is not None:
                    piece_x, piece_y = get_piece_coord(i)
                    piece_x = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.cos(angle_step*(5*piece_x+piece_y)/5))
                    piece_y = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.sin(angle_step*(5*piece_x+piece_y)/5))
                    rect = pygame.Surface((20, 20)).get_rect(center=(piece_x, piece_y))
                    if rect.collidepoint(x, y):
                        dest_row, dest_col = get_piece_coord(i)
                        break

            moves = get_available_moves(selected_piece, pieces)
            for move in moves:
                jumper, target, destination = move
                if jumper == selected_piece and destination == 5*dest_row + dest_col:
                    pieces[jumper] = None
                    pieces[target] = None
                    pieces[destination] = pygame.Surface((20, 20))
                    pieces[destination].fill(RED)
                    pieces[destination].set_colorkey(WHITE)
                    pygame.draw.circle(pieces[destination], BLACK, (10, 10), 10, 2)
                    # switch turn
                    player_turn = not player_turn
                    selected_piece = None
                    break

    # Let AI move
    if not player_turn:
        piece, move = ai_move(pieces)
        if piece is not None and move is not None:
            jumper, target, destination = move
            dest_row, dest_col = get_piece_coord(destination)
            dest_x = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.cos(angle_step*(5*dest_row+dest_col)/5))
            dest_y = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.sin(angle_step*(5*dest_row+dest_col)/5))
            jumper_x, jumper_y = get_piece_coord(jumper)
            jumper_x = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.cos(angle_step*(5*jumper_x+jumper_y)/5))
            jumper_y = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.sin(angle_step*(5*jumper_x+jumper_y)/5))
            target_x, target_y = get_piece_coord(target)
            target_x = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.cos(angle_step*(5*target_x+target_y)/5))
            target_y = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.sin(angle_step*(5*target_x+target_y)/5))

            pieces[jumper] = None
            pieces[target] = None
            pieces[destination] = pygame.Surface((20, 20))
            pieces[destination].fill(RED)
            pieces[destination].set_colorkey(WHITE)
            pygame.draw.circle(pieces[destination], BLACK, (10, 10), 10, 2)
            pygame.draw.line(window, BLACK, (jumper_x, jumper_y), (dest_x, dest_y), 2)
            pygame.draw.circle(window, WHITE, (target_x, target_y), 10)
            # switch turn
            player_turn = not player_turn

    # draw pieces
    for i in range(len(pieces)):
        if pieces[i] is not None:
            row, col = get_piece_coord(i)
            x = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.cos(angle_step*(5*row+col)/5))
            y = int(BOARD_RADIUS + (BOARD_RADIUS-3*space)*math.sin(angle_step*(5*row+col)/5))
            rect = pygame.Rect(x-10, y-10, 20, 20)
            window.blit(pieces[i], rect)

    if player_turn:
        pygame.display.set_caption('Jumping Game: Player Turn')
    else:
        pygame.display.set_caption('Jumping Game: AI Turn')

    pygame.display.update()
    fps_clock.tick(FPS)

